from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.pool import QueuePool
from typing import Optional, List, Dict, Any
from datetime import datetime, timedelta
import pandas as pd

from .models import Base, MarketData, OrderBook, Trade, Position, TradingSignal

class Database:
    """数据库管理类"""
    
    def __init__(self, db_url: str = 'sqlite:///crypto_robot.db'):
        """
        初始化数据库连接
        
        Args:
            db_url: 数据库连接URL
        """
        self.engine = create_engine(
            db_url,
            poolclass=QueuePool,
            pool_size=5,
            max_overflow=10,
            pool_timeout=30
        )
        self.Session = scoped_session(sessionmaker(bind=self.engine))
        
    def init_db(self):
        """初始化数据库表"""
        Base.metadata.create_all(self.engine)
        
    def save_market_data(self, data: Dict[str, Any]):
        """
        保存市场数据
        
        Args:
            data: 市场数据字典
        """
        session = self.Session()
        try:
            market_data = MarketData(**data)
            session.add(market_data)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def save_orderbook(self, data: Dict[str, Any]):
        """
        保存订单簿数据
        
        Args:
            data: 订单簿数据字典
        """
        session = self.Session()
        try:
            orderbook = OrderBook(**data)
            session.add(orderbook)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def save_trade(self, data: Dict[str, Any]):
        """
        保存交易记录
        
        Args:
            data: 交易记录字典
        """
        session = self.Session()
        try:
            trade = Trade(**data)
            session.add(trade)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def save_position(self, data: Dict[str, Any]):
        """
        保存持仓记录
        
        Args:
            data: 持仓记录字典
        """
        session = self.Session()
        try:
            position = Position(**data)
            session.add(position)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def save_trading_signal(self, data: Dict[str, Any]):
        """
        保存交易信号
        
        Args:
            data: 交易信号字典
        """
        session = self.Session()
        try:
            signal = TradingSignal(**data)
            session.add(signal)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def get_market_data(self, exchange: str, symbol: str,
                       start_time: datetime, end_time: datetime) -> pd.DataFrame:
        """
        获取市场数据
        
        Args:
            exchange: 交易所名称
            symbol: 交易对
            start_time: 开始时间
            end_time: 结束时间
            
        Returns:
            DataFrame格式的市场数据
        """
        session = self.Session()
        try:
            query = session.query(MarketData).filter(
                MarketData.exchange == exchange,
                MarketData.symbol == symbol,
                MarketData.timestamp.between(start_time, end_time)
            ).order_by(MarketData.timestamp.asc())
            
            data = [{
                'timestamp': item.timestamp,
                'open': item.open,
                'high': item.high,
                'low': item.low,
                'close': item.close,
                'volume': item.volume
            } for item in query.all()]
            
            return pd.DataFrame(data)
        finally:
            session.close()
            
    def get_latest_trades(self, exchange: str, symbol: str, limit: int = 100) -> List[Dict]:
        """
        获取最近的交易记录
        
        Args:
            exchange: 交易所名称
            symbol: 交易对
            limit: 返回记录数量
            
        Returns:
            交易记录列表
        """
        session = self.Session()
        try:
            trades = session.query(Trade).filter(
                Trade.exchange == exchange,
                Trade.symbol == symbol
            ).order_by(Trade.timestamp.desc()).limit(limit).all()
            
            return [{
                'trade_id': trade.trade_id,
                'order_id': trade.order_id,
                'timestamp': trade.timestamp,
                'side': trade.side,
                'type': trade.type,
                'price': trade.price,
                'amount': trade.amount,
                'fee': trade.fee,
                'status': trade.status
            } for trade in trades]
        finally:
            session.close()
            
    def get_open_positions(self, exchange: str = None) -> List[Dict]:
        """
        获取当前持仓
        
        Args:
            exchange: 交易所名称（可选）
            
        Returns:
            持仓记录列表
        """
        session = self.Session()
        try:
            query = session.query(Position).filter(Position.status == 'open')
            if exchange:
                query = query.filter(Position.exchange == exchange)
                
            positions = query.all()
            return [{
                'exchange': pos.exchange,
                'symbol': pos.symbol,
                'side': pos.side,
                'entry_price': pos.entry_price,
                'current_price': pos.current_price,
                'amount': pos.amount,
                'leverage': pos.leverage,
                'unrealized_pnl': pos.unrealized_pnl,
                'realized_pnl': pos.realized_pnl
            } for pos in positions]
        finally:
            session.close()
            
    def get_pending_signals(self, exchange: str = None) -> List[Dict]:
        """
        获取未执行的交易信号
        
        Args:
            exchange: 交易所名称（可选）
            
        Returns:
            未执行的交易信号列表
        """
        session = self.Session()
        try:
            query = session.query(TradingSignal).filter(TradingSignal.executed == False)
            if exchange:
                query = query.filter(TradingSignal.exchange == exchange)
                
            signals = query.all()
            return [{
                'exchange': signal.exchange,
                'symbol': signal.symbol,
                'signal_type': signal.signal_type,
                'side': signal.side,
                'price': signal.price,
                'strength': signal.strength,
                'timestamp': signal.timestamp
            } for signal in signals]
        finally:
            session.close()
            
    def mark_signal_executed(self, signal_id: int):
        """
        标记交易信号为已执行
        
        Args:
            signal_id: 信号ID
        """
        session = self.Session()
        try:
            signal = session.query(TradingSignal).get(signal_id)
            if signal:
                signal.executed = True
                session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
            
    def cleanup_old_data(self, days: int = 30):
        """
        清理旧数据
        
        Args:
            days: 保留最近多少天的数据
        """
        session = self.Session()
        try:
            cutoff_date = datetime.utcnow() - timedelta(days=days)
            
            # 清理市场数据
            session.query(MarketData).filter(
                MarketData.timestamp < cutoff_date
            ).delete()
            
            # 清理订单簿数据
            session.query(OrderBook).filter(
                OrderBook.timestamp < cutoff_date
            ).delete()
            
            # 清理已完成的交易记录
            session.query(Trade).filter(
                Trade.timestamp < cutoff_date,
                Trade.status.in_(['filled', 'canceled'])
            ).delete()
            
            # 清理已关闭的持仓记录
            session.query(Position).filter(
                Position.updated_at < cutoff_date,
                Position.status == 'closed'
            ).delete()
            
            # 清理已执行的交易信号
            session.query(TradingSignal).filter(
                TradingSignal.timestamp < cutoff_date,
                TradingSignal.executed == True
            ).delete()
            
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
